package cn.com.cetccst.tools.encrypt;

import android.util.Base64;

import java.io.ByteArrayInputStream;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class Encryption {

    // In message transform v1, id size is 2x64b for connId + 32b for msgId.
    // The version code is "1". Using byte array so we can feed into cipher.
    public static final int MESSAGE_ID_SIZE_V1 = 20;
    public static final byte[] MESSAGE_TRANSFORM_ID_V1 = {1};
    public static final int MESSAGE_IV_SIZE = 16;
    private final static String HEX = "0123456789abcdef";
    public static Base64Util base64 = new Base64Util();
    private static int NUM_CYCLES = 20011;

    public static byte[] decrypt(String cipherText, String password, String salt) {
        // Use key + salt to decrypt the decoded data
        try {
            // Decode the data from Base 64
            byte cipherData[] = base64.decode(cipherText.getBytes(), Base64.NO_WRAP);

            Cipher pbeCipher = init(password, salt, Cipher.DECRYPT_MODE);

            // Encrypt the cleartext
            byte[] plainText = pbeCipher.doFinal(cipherData);

            return (plainText);

        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    public static String encrypt(byte cleartext[], String password, String salt) {
        // Use key + salt to decrypt the decoded data
        try {
            Cipher pbeCipher = init(password, salt, Cipher.ENCRYPT_MODE);

            // Encrypt the cleartext
            byte[] ciphertext = pbeCipher.doFinal(cleartext);

            // Encode the data to Base 64
            return (new String(base64.encode(ciphertext, Base64.NO_WRAP)));

        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    public static byte[] decode(String encodedtext) {
        try {
            return base64.decode(encodedtext.getBytes(), Base64.NO_WRAP);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String encode(byte cleartext[]) {
        try {
            return (new String(base64.encode(cleartext, Base64.NO_WRAP)));
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String generateInitVector() {
        SecureRandom random = new SecureRandom();
        byte bytes[] = new byte[8];
        random.nextBytes(bytes);

        return new String(base64.encode(bytes, Base64.DEFAULT));
    }

    public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        // Generate a pub/priv key pair
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        SecureRandom random = new SecureRandom();
        keyGen.initialize(2048, random);
        return keyGen.generateKeyPair();
    }

    public static X509Certificate extractCertificate(String encodedCert)
            throws CertificateException {
        if (encodedCert == null)
            throw new CertificateException("extractCertificate: Certificate is null");

        ByteArrayInputStream bis =
                new ByteArrayInputStream(base64.decode(encodedCert.getBytes(), Base64.NO_WRAP));
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate cert = (X509Certificate) cf.generateCertificate(bis);

        return cert;
    }

    public static SecretKey getHardcodedKey() {
        byte[] hardcodedKey = {1, 2, 3, 4, 5, 6, 7, 8, 9,
                10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
                20, 21, 22, 23, 24};
        SecretKeySpec key = new SecretKeySpec(hardcodedKey, 0, 24, "AES");
        return key;
    }

    public static SecretKey extractSecretKey(String encodedKey,
                                             PrivateKey unwrappingKey)
            throws InvalidKeyException {
        if (encodedKey == null) {
            throw new InvalidKeyException("extractSecretKey: null enc key");
        }
        if (unwrappingKey == null) {
            throw new InvalidKeyException("extractSecretKey: null unwrapper");
        }
        byte[] encryptedBytes = base64.decode(encodedKey.getBytes(),
                Base64.NO_WRAP);

        try {
            // Decrypt the symmetric key.
            Cipher unwrapCipher =
                    Cipher.getInstance("RSA/None/OAEPWithSHA1AndMGF1Padding",
                            "BC");
            if (unwrapCipher == null) {
                throw new InvalidKeyException("extractSecretKey: null cipher");
            }
            unwrapCipher.init(Cipher.DECRYPT_MODE,
                    unwrappingKey,
                    new SecureRandom());
            byte[] decryptedBytes = unwrapCipher.doFinal(encryptedBytes);

            // We will only use the first bytes in what we got.
            SecretKeySpec key = new SecretKeySpec(decryptedBytes, 0, 24,
                    "AES");
            return key;
        } catch (Exception e) {
            e.printStackTrace();
            throw new InvalidKeyException("could not extract secret key");
        }
    }

    public static String generatePassword() {
        SecureRandom random = new SecureRandom();
        byte pwd[] = new byte[128];
        random.nextBytes(pwd);

        return new String(base64.encode(pwd, Base64.DEFAULT));
    }

    // Construct a 20-byte message ID from a few building blocks:
    // a message counter, a source connection ID, and target connection ID.
    // Note that here we don't control how these IDs are obtained.
    // For example, all IDs might be dictated by the initiator, in which
    // case the target is vulnerable to replay attacks by a MITM.
    public static byte[] messageComposeId(int ctr, long nonce1, long nonce2) {
        byte[] result = new byte[MESSAGE_ID_SIZE_V1];
        storeLongAsBytes(result, 0, (long) ctr, 4);
        storeLongAsBytes(result, 4, nonce1, 8);
        storeLongAsBytes(result, 12, nonce2, 8);
        return result;
    }

    // Encrypt. Make sure to use a fresh IV each time. Attach a MAC.
    //
    // Wrapped message structure (in bytes):
    //   0   .. 15           IV
    //   16  .. X            Ek({1}|msg)
    //   X+1 .. X+hMacLength Hk(msgId|Ek({1}|msg))
    //
    public static byte[] messageEncryptThenAuthenticate(byte[] msg,
                                                        byte[] msgId,
                                                        SecretKey key)
            throws InvalidKeyException {
        if (msgId.length != MESSAGE_ID_SIZE_V1) {
            throw new InvalidKeyException("messageEtA: id size mismatch");
        }
        byte[] encMsg;
        try {
            SecureRandom rnd = new SecureRandom();
            byte[] ivBytes = new byte[MESSAGE_IV_SIZE];
            rnd.nextBytes(ivBytes);
            IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);

            Mac hMac = Mac.getInstance("HmacSHA256", "BC");
            Key hMacKey = new SecretKeySpec(key.getEncoded(), "HmacSHA256");
            if (hMac == null || hMacKey == null) {
                throw new InvalidKeyException("messageEtA: null Hmac");
            }

            Cipher encCipher =
                    Cipher.getInstance("AES/CBC/PKCS5Padding", "BC");
            if (encCipher == null) {
                throw new InvalidKeyException("messageEtA: null cipher");
            }
            encCipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
            encMsg = new byte[ivBytes.length +
                    encCipher.getOutputSize(1 + msg.length) +
                    hMac.getMacLength()];
            System.arraycopy(ivBytes, 0, encMsg, 0, ivBytes.length);
            // Make sure to start the wrapped message with the version id,
            // so that future versions can be backwards compatible.
            int cLen = encCipher.update(MESSAGE_TRANSFORM_ID_V1, 0, 1,
                    encMsg, ivBytes.length);
            cLen += encCipher.doFinal(msg, 0, msg.length,
                    encMsg, ivBytes.length + cLen);

            hMac.init(hMacKey);
            hMac.update(msgId);
            hMac.update(encMsg, ivBytes.length, cLen);
            if (cLen != (encMsg.length - ivBytes.length - hMac.getMacLength())) {
                throw new InvalidKeyException("messageEtA: length error");
            }
            hMac.doFinal(encMsg, ivBytes.length + cLen);
        } catch (Exception e) {
            e.printStackTrace();
            throw new InvalidKeyException("could not encrypt message");
        }

        return encMsg;
    }

    // Check the MAC. Decrypt.
    public static byte[] messageAuthenticateThenDecrypt(byte[] msg,
                                                        byte[] msgId,
                                                        SecretKey key)
            throws InvalidKeyException {
        if (msgId.length != MESSAGE_ID_SIZE_V1) {
            throw new InvalidKeyException("messageDaA: id size mismatch");
        }
        byte[] decMsg;
        try {
            IvParameterSpec ivSpec =
                    new IvParameterSpec(msg, 0, MESSAGE_IV_SIZE);

            Mac hMac = Mac.getInstance("HmacSHA256", "BC");
            Key hMacKey = new SecretKeySpec(key.getEncoded(), "HmacSHA256");
            if (hMac == null || hMacKey == null) {
                throw new InvalidKeyException("messageDaA: null Hmac");
            }
            if (msg.length < (hMac.getMacLength() + MESSAGE_IV_SIZE)) {
                throw new InvalidKeyException("messageDaA: message too short");
            }

            Cipher decCipher =
                    Cipher.getInstance("AES/CBC/PKCS5Padding", "BC");
            if (decCipher == null) {
                throw new InvalidKeyException("messageDaA: null cipher");
            }
            decCipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
            decMsg = decCipher.doFinal(msg, MESSAGE_IV_SIZE,
                    msg.length
                            - MESSAGE_IV_SIZE
                            - hMac.getMacLength());
            if (decMsg.length < 1) {
                throw new InvalidKeyException("messageDaA: short message");
            }
            // Check version id.
            if (decMsg[0] != MESSAGE_TRANSFORM_ID_V1[0]) {
                throw new InvalidKeyException("messageDaA: bad version");
            }

            hMac.init(hMacKey);
            hMac.update(msgId);
            hMac.update(msg, MESSAGE_IV_SIZE,
                    msg.length - MESSAGE_IV_SIZE - hMac.getMacLength());
            byte[] hMacResult = hMac.doFinal();
            byte[] hMacFromMsg = new byte[hMac.getMacLength()];
            System.arraycopy(msg, msg.length - hMac.getMacLength(),
                    hMacFromMsg, 0, hMac.getMacLength());
            if (!MessageDigest.isEqual(hMacResult, hMacFromMsg)) {
                throw new InvalidKeyException("messageDaA: HMAC mismatch");
            }
        } catch (Exception e) {
            e.printStackTrace();
            throw new InvalidKeyException("could not decrypt message");
        }

        byte[] result = new byte[decMsg.length - 1];
        System.arraycopy(decMsg, 1, result, 0, decMsg.length - 1);
        return result;
    }

    public static String toHex(byte[] buf) {
        if (buf == null)
            return "";
        StringBuffer result = new StringBuffer(2 * buf.length);
        for (int i = 0; i < buf.length; i++) {
            appendHex(result, buf[i]);
        }
        return result.toString();
    }

    private static Cipher init(String password, String salt, int mode) throws Exception {
        // Process password and salt
        PBEParameterSpec pbeParamSpec = new PBEParameterSpec(
                base64.decode(salt.getBytes(), Base64.DEFAULT), NUM_CYCLES);
        PBEKeySpec pbeKeySpec = new PBEKeySpec(password.toCharArray());
        SecretKeyFactory keyFac =
                SecretKeyFactory.getInstance("PBEWITHSHA256AND128BITAES-CBC-BC",
                        "BC");
        SecretKey pbeKey = keyFac.generateSecret(pbeKeySpec);

        // Create Cipher
        Cipher pbeCipher =
                Cipher.getInstance("PBEWITHSHA256AND128BITAES-CBC-BC",
                        "BC");
        pbeCipher.init(mode, pbeKey, pbeParamSpec);

        return pbeCipher;
    }

    private static void storeLongAsBytes(byte[] buf, int offs,
                                         long value, int size) {
        for (int ix = 0; ix < size; ix++) {
            buf[offs + ix] = (byte) value;
            value = value >>> 8;
        }
    }

    private static void appendHex(StringBuffer sb, byte b) {
        sb.append(HEX.charAt((b >> 4) & 0x0f)).append(HEX.charAt(b & 0x0f));
    }

    // Stub to help with unit testing
    public static class Base64Util {
        public byte[] decode(byte[] bytes, int flags) {
            return Base64.decode(bytes, flags);
        }

        public byte[] encode(byte[] bytes, int flags) {
            return Base64.encode(bytes, flags);
        }
    }
}
